#!/usr/bin/env python
# encoding: utf-8
"""
primer3_tag.py

Created by Brant Faircloth on 2006-10-01.
Copyright (c) 2006-2007 Brant C. Faircloth. All rights reserved.

Released under the GNU General Public License 2.0 
http://www.gnu.org/licenses/gpl.html

"""


import sys, os, string, csv, pdb

class mods: 
    """class defining several methods we will use to modify dna sequence.""" 
 
    def __init__(self, s): 
        """create DNA instance initialized to string s.""" 
        self.seq = s 
        
    def complement(self):
        """returns complementary dna sequence"""
        #--------------------------------------------------------------------------------------
        # first, we capitalize the sequence (if lowercase), then we translate it to its reverse
        # complement.
        #--------------------------------------------------------------------------------------
        if self.seq.islower():self.seq = self.seq.upper() 
        tab = string.maketrans('AGCTN','TCGAN')
        output = string.translate(self.seq, tab)
        return output

    def reverseComplement(self):
        """returns reverse complementary dna sequence"""
        #-----------------------------------------------------------------------------------------
        # the self.seq (sequence instance) must be converted to a list, reversed, put back together
        # as a string, and then complemented.
        #-----------------------------------------------------------------------------------------
        listS = list(self.seq)
        listS.reverse()
        self.seq = (''.join(listS))
        return self.complement()
        
class tagMain:
    """main class for sequence building and input file construction"""
    #--------------------------------------------------------------------------------------
    # This class is the main class we will use for building the temporary outfile for use
    # with primer3.  It is subclassed by several other classes and that allows us to only
    # define the main function 1X, thus saving us some code.
    #--------------------------------------------------------------------------------------
    def __init__(self, tag, upper, lower):
        self.fragTag=tag
        #---------------------------------------------------------------------------------
        # The vector sequence used here is the first 250 bp of Enterobacteria phage lambda
        # GenBank gi:9626243.  It really serves no other purpose than a spacer to put
        # between the primers we will be testing.
        #---------------------------------------------------------------------------------
        self.lamVector='GGGCGGCGACCTCGCGGGTTTTCGCTATTTATGAAAATTTTCCGGTTTAAGGCGTTTCCGTTCTTCTTCGTCATAACTTAATGTTTTTATTTAAAATACCCTCTGAAAAGAAAGGAAACGACAGGTGCTGAAAGCGAGGCTTTTTGGCCTCTGTCGTTTCCTTTCTCTGTTTTTGTCCGTGGAATGAACAATGGAAGTCAACAAAAAGCAGCTGGCTGACATTTTCGGTGCGAGTATCCGTACCATTCAG'
        self.upperP = upper
        self.lowerP = lower
        self.lowerPRevComp = mods(self.lowerP).reverseComplement()
        #-------------------------------------------------------------------------------------
        # The above is essentially setting the variables which will be used in the subclasses.
        # That way, we don't have to define them again for each new function
        #-------------------------------------------------------------------------------------
        
    def commonBases(self, tag, primer):
        """returns tag; checks primer and tag for common bases and removes them"""
        #global common
        self.common = False
        #print 'input tag', tag                            #this crap can go
        i=5
        while i > 0:
            if tag.endswith(primer[0:i]):
                tag = tag.rstrip(primer[0:i])
                self.common = True
                #print 'tag was modified'                           #this crap can go
                #print 'ouput tag', tag                    #this crap can go
                break
            else:
                i -= 1
        return tag 
    
    def writeFile(self):
        """function builds temporary output file to be used by primer 3, subclassed below"""
        self.outFile=open('primer3_tag_temp.txt','w')
        self.outFile.write(('SEQUENCE=%s\n') % (self.seqBuild()))
        if self.common:
            self.outFile.write('PRIMER_COMMENT=Common Bases Modified\n')
        self.tagAdd()
        #---------------------------------------------------------------------------------
        # Set (or unset) many default values/weights for primer3.  Range must be set high, 
        # TMs are set ridiculously high as the tag adds a lot of temp.
        #---------------------------------------------------------------------------------
        self.outFile.write('PRIMER_PRODUCT_SIZE_RANGE=100-500\n')
        self.outFile.write('PRIMER_MIN_TM=45.0\n')
        self.outFile.write('PRIMER_MAX_TM=100.0\n')
        self.outFile.write('PRIMER_SELF_ANY=6.0\n')
        self.outFile.write('PRIMER_SELF_END=2.5\n')
        self.outFile.write('PRIMER_MIN_SIZE=15\n')
        self.outFile.write('PRIMER_MAX_SIZE=45\n')
        self.outFile.write('PRIMER_PICK_ANYWAY=1\n')
        self.outFile.write('PRIMER_EXPLAIN_FLAG=1\n')
        self.outFile.write('PRIMER_WT_TM_GT=0.0\n')
        self.outFile.write('PRIMER_WT_TM_LT=0.0\n')
        self.outFile.write('PRIMER_WT_SIZE_GT=0.0\n')
        self.outFile.write('PRIMER_WT_SIZE_LT=0.0\n')
        self.outFile.write('PRIMER_WT_COMPL_ANY=1.0\n')
        self.outFile.write('PRIMER_WT_COMPL_END=1.0\n')
        self.outFile.write('PRIMER_WT_END_STABILITY=1.0\n')
        self.outFile.write('PRIMER_PAIR_WT_COMPL_ANY=1.0\n')
        self.outFile.write('PRIMER_PAIR_WT_COMPL_END=1.0\n')
        self.outFile.write('PRIMER_WT_GC_PERCENT_LT=0.0\n')   #do we want to keep GC?
        self.outFile.write('PRIMER_WT_GC_PERCENT_LT=0.0\n')   #do we want to keep GC? (do we want to define GC?)
        self.outFile.write('=')
        self.outFile.close()
        #------------------------
        # End primer3 file values
        #------------------------

class tagFront(tagMain):
    """subclass of tagMain; front tag (forward/upper primer) sequence building and input file creation"""
    
    def seqBuild(self):
        """builds contig for primer selection"""
        self.fragTag = self.commonBases(self.fragTag, self.upperP)
        return self.fragTag + self.upperP + self.lamVector + self.lowerPRevComp
        #------------------------------------------------------------------------------------------------------
        # we are 'building' a contig for primer3 that contains the tag we want to use, the upper primer we are 
        # testing, the vector dna, and the reverse complement of our lower primer.
        #------------------------------------------------------------------------------------------------------
    
    def tagAdd(self):
        """builds upper primer with tag and sets lower primer value"""
        self.outFile.write(('PRIMER_LEFT_INPUT=%s\n') % (self.fragTag + self.upperP))
        self.outFile.write(('PRIMER_RIGHT_INPUT=%s\n') % (self.lowerP))
        return 

class tagRear(tagMain):
    """subclass of tagMain; rear tag (reverse/lower primer) sequence building and input file creation"""
    
    def seqBuild(self):
        """builds contig for primer selection"""
        self.fragTag = self.commonBases(self.fragTag, self.lowerP)
        self.fragTagRevComp = mods(self.fragTag).reverseComplement()
        #----------------------------------------------------------------------------
        # the above is required so we can generate the reverseComp of fragTag *after*
        # we remove common bases (before would hose us)
        #----------------------------------------------------------------------------
        return self.upperP + self.lamVector + self.lowerPRevComp + self.fragTagRevComp
    
    def tagAdd(self):
        """builds lower primer with tag and sets upper primer value"""
        self.outFile.write(('PRIMER_LEFT_INPUT=%s\n') % (self.upperP))
        self.outFile.write(('PRIMER_RIGHT_INPUT=%s\n') % (self.fragTag + self.lowerP))
        return

class primerInfo:
    """this class will test out tagged primers"""
    
    def primer3Run(self):
        pathToPrimer3 = '~/bin/primer3_core < primer3_tag_temp.txt'
        #pathToPrimer3 = '../src/primer3_core < ../scripts/temp.txt'
        primer3Output=string.split(os.popen(pathToPrimer3).read(),'\n')
        #--------------------------------------------------------------------------------------
        # here we call primer3, read the stdout from primer3, and put that output into a
        # string.  The chop up the string and put it in a dictionary - ahhhhh, dictionaries.
        #--------------------------------------------------------------------------------------
        primerDataDict={}
        for item in primer3Output[0:len(primer3Output)-2]:      #drop last 2 lines - this is an '=' and a space
            splitItem=string.split(item,'=')
            try:
                splitItem[1]=float(splitItem[1])                #convert what we can to floats
            except:
                pass
            primerDataDict[splitItem[0]]=splitItem[1]     
        stuffToKeep=['PRIMER_LEFT_END_STABILITY', 'PRIMER_LEFT_INPUT', 'PRIMER_LEFT_PENALTY', 'PRIMER_LEFT_SELF_ANY','PRIMER_LEFT_SELF_END', 'PRIMER_PAIR_COMPL_ANY', 'PRIMER_PAIR_COMPL_END', 'PRIMER_PAIR_PENALTY', 'PRIMER_RIGHT_END_STABILITY', 'PRIMER_RIGHT_INPUT', 'PRIMER_RIGHT_PENALTY', 'PRIMER_RIGHT_SELF_ANY', 'PRIMER_RIGHT_SELF_END', 'PRIMER_WARNING','PRIMER_COMMENT']
        for item in primerDataDict.keys():
            if item not in stuffToKeep:
                try:
                    del primerDataDict[item]
                except:
                    pass
        return primerDataDict

def getData(dataFile):
    """get the freaking data"""
    dataDialect = csv.Sniffer().sniff(dataFile)                 # try to determine dialect automagically
    headerTrue = csv.Sniffer().has_header(dataFile)             # try to determine if file has header row
    primerData = open(dataFile,'r')                               # open datafile for read
    primerReader = csv.reader(primerData, delimiter=',', dialect = dataDialect)  # setup csv.reader
    if headerTrue:
        primerReader.next()
    primerDict = {}
    for row in primerReader:
        clone,left,right = row[0],row[1],row[4]
        primerDict[clone]=[left,right]
    primerData.close()
    return primerDict

def writeData(clone, primer, header, output):
    if 'PRIMER_PAIR_PENALTY' in primer.keys():
        #pdb.set_trace()
        buildList = [clone]
        for element in header[1:]:
            try:
                try:
                    elementFloat = float(primer[element])
                except:
                    elementFloat = primer[element]
                buildList.append(elementFloat)
            except:
                buildList.append('')
        #if common:
        #    buildList.append('Common bases modified')
        else: buildList.append('')
        print buildList
        output.writerow(buildList)
    else: pass

def main():
    cagTag = 'CAGTCGGGCGTCATCA'
    m13rTag = 'GGAAACAGCTATGACCAT'
    primer3data = 'primer3/primerNoTag.csv'
    data = getData(primer3data)
    outputFile = 'primer3/primerTag.csv'
    outputData = open(outputFile,'w')                              # open datafile for read
    outputWriter = csv.writer(outputData, dialect = 'excel')
    header = [
    'Clone',
    'PRIMER_LEFT_INPUT',
    'PRIMER_LEFT_SELF_ANY',
    'PRIMER_LEFT_SELF_END',
    'PRIMER_RIGHT_INPUT',
    'PRIMER_RIGHT_SELF_ANY',
    'PRIMER_RIGHT_SELF_END',
    'PRIMER_PAIR_COMPL_ANY',
    'PRIMER_PAIR_COMPL_END',
    'PRIMER_WARNING',
    'PRIMER_COMMENT']
    outputWriter.writerow(header)
    print data

    for item in data:
        if data[item][0] == '' or data[item][1]=='':
            pass
        else:
            #print item[0]
            resultsDict={}
            #print item[0], item[1], item[2]
            tagFront(cagTag, data[item][0], data[item][1]).writeFile()
            resultsDict[0]=primerInfo().primer3Run()
            tagRear(cagTag, data[item][0], data[item][1]).writeFile()
            resultsDict[1]=primerInfo().primer3Run()
            tagFront(m13rTag, data[item][0], data[item][1]).writeFile()
            resultsDict[2]=primerInfo().primer3Run()
            tagRear(m13rTag, data[item][0], data[item][1]).writeFile()
            resultsDict[3]=primerInfo().primer3Run()
            i=0; primerPicks=[]
            #print resultsDict
            for element in resultsDict:
                try:
                    primerPicks.append(resultsDict[element]['PRIMER_PAIR_PENALTY'])
                    i+=1                                               
                except:
                    primerPicks.append(1000)  #insert arbitrarily bad value to list for primer w/ no results
                    i+=1
            try:
                goodPrimer = resultsDict[primerPicks.index(min(primerPicks))]        #choose MINIMUM primer penalty value and keep that record     
                del resultsDict[primerPicks.index(min(primerPicks))]
                writeData(item,goodPrimer,header,outputWriter)
            except:
                pass
            try:
                sortedGoodPrimer = goodPrimer.keys()
                sortedGoodPrimer.sort()     
                print '-------------------------------------'
                print (("--- Ideal Primer for Clone %s---") % (item))
                print '-------------------------------------\n'
                #for character in sorted(goodPrimer):
                for character in sortedGoodPrimer:
                    print (('%s = %s') % (character, goodPrimer[character]))             
                print '-----------------------------------\n'
                print '\t---------------------------------------'
                print(("\t--- Other Primers for Clone %s ---") % (item))
                print '\t---------------------------------------'
                for foo in resultsDict:
                    print (("\t--- Other Primer %s---") % (foo)) 
                    sortedOtherResults = resultsDict[foo].keys()
                    sortedOtherResults.sort()
                    for bar in sortedOtherResults:
                    #for bar in resultsDict[foo]:
                    #for bar in sorted(resultsDict[foo]):
                        print (('\t%s = %s') % (bar, resultsDict[foo][bar])) 
                    print '\n'
            except:
                pass
    outputData.close()
    os.remove('primer3_tag_temp.txt')

if __name__ == '__main__':
    main()                                       